from ocr.utils.tools import read_image
from ocr.utils.results import ServiceResult


class MainStreamController:
    def __init__(self, model_manager):
        self.document_image_classifier = model_manager.document_images_classifier

    def process(self, inputs):
        res = ServiceResult()

        img = read_image(inputs)
        if img is None:
            res = res.set_error_message("read image failed")
            return res

        image_class = self.document_image_classifier.predict(img)
        if image_class is None:
            res = res.set_error_message("predict images failed")
            return res

        res.set_success_message(image_class)
        return res
